package auth

import (
	"context"
	"crypto/rsa"
	"gitee.com/jacktuotuo/coolcar/pkg/auth/token"
	"github.com/go-kratos/kratos/v2/log"
	"github.com/go-kratos/kratos/v2/middleware"
	"github.com/go-kratos/kratos/v2/transport/http"
	"github.com/golang-jwt/jwt/v4"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	"io/ioutil"
	"os"
	"strings"
)

type interceptor struct {
	verifier verifier
}

type verifier interface {
	Verify(string) (string, error)
}

type accountIDKey struct{}

func AccountIdFromContext(ctx context.Context) (AccountId, error) {
	v := ctx.Value(accountIDKey{})

	aid, ok := v.(AccountId)

	if !ok {
		return "", status.Error(codes.Unauthenticated, "")
	}

	return aid, nil
}

func ContextWithAid(ctx context.Context, aid AccountId) context.Context {
	return context.WithValue(ctx, accountIDKey{}, aid)
}

const (
	authorizationHeader = "Authorization"
	bearerPrefix        = "Bearer "
)

type AccountId string

func (a AccountId) String() string {
	return string(a)
}

func tokenFromContext(ctx context.Context) (string, error) {

	Unauthenticated := status.Error(codes.Unauthenticated, "")

	tkn := ""

	wrapper, ok := ctx.(http.Context)

	if !ok {
		return "", Unauthenticated
	}

	m := wrapper.Header()

	for _, v := range m[authorizationHeader] {
		if strings.HasPrefix(v, bearerPrefix) {
			tkn = v[len(bearerPrefix):]
		}
	}

	if tkn == "" {
		return "", Unauthenticated
	}

	return tkn, nil

}

// Auth is  auth middleware.
func Auth(logger log.Logger, publicKeyFile string) middleware.Middleware {
	var publicKey *rsa.PublicKey

	l := log.NewHelper(logger)

	f, err := os.Open(publicKeyFile)
	if err != nil {
		l.Errorf("can not open publicKey file :%v", err)
	} else {
		b, err := ioutil.ReadAll(f)
		if err != nil {
			l.Errorf("can not read publicKey file :%v", err)
		} else {
			publicKey, err = jwt.ParseRSAPublicKeyFromPEM(b)
			if err != nil {
				l.Errorf("can not parse publicKey :%v", err)
			}
		}
	}

	i := interceptor{
		verifier: &token.JwtTokenVerifier{
			PublicKey: publicKey,
		},
	}

	return func(handler middleware.Handler) middleware.Handler {
		return func(ctx context.Context, req interface{}) (reply interface{}, err error) {

			tkn, err := tokenFromContext(ctx)
			if err != nil {
				return nil, status.Error(codes.Unauthenticated, "")
			}

			aid, err := i.verifier.Verify(tkn)
			if err != nil {
				return nil, status.Errorf(codes.Unauthenticated, "token not valid :%v", err)
			}

			ctx = ContextWithAid(ctx, AccountId(aid))
			return handler(ctx, req)
		}
	}
}
